#include <boost/make_shared.hpp>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/point_representation.h>

#include <pcl/io/pcd_io.h>

#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/filter.h>

#include <pcl/features/normal_3d.h>

#include <pcl/registration/icp.h>
#include <pcl/registration/icp_nl.h>
#include <pcl/registration/transforms.h>

#include <pcl/visualization/pcl_visualizer.h>

using pcl::visualization::PointCloudColorHandlerCustom;
using pcl::visualization::PointCloudColorHandlerGenericField;

// convenient typedefs
typedef pcl::PointXYZ PointT;
typedef pcl::PointCloud<PointT> PointCloud;
typedef pcl::PointNormal PointNormalT;
typedef pcl::PointCloud<PointNormalT> PointCloudWithNormals;

/**
 * @brief params for registration
 * 
 */
double VoxelGridLeafSize(0.05);  // default: 0.05
int NormEstKSearch(25);         // default: 30
double ICPMaxDist(0.25);         // default: 0.1
int ICPMaxInnerIteration(2);    // default: 2
int ICPMaxExternIteration(90);  // default: 30
double ICPEPS(1e-09);           // default: 1e-06


/**
 * @brief visualizer related
 * 
 */
// visualizer
pcl::visualization::PCLVisualizer *p;
// its left and right viewports
int vp_1, vp_2;

/**
 * @brief custom pointcloud struct
 * 
 */
struct PCD
{
    PointCloud::Ptr cloud;
    std::string f_name;

    // constructor
    PCD() : cloud(new PointCloud){};
};

/**
 * @brief struct to compare PCD f_name
 * 
 */
struct PCDComparator
{
    bool operator()(const PCD &p1, const PCD &p2)
    {
        return (p1.f_name < p2.f_name);
    }
};

/**
 * @brief class MyPointRepresentation
 * Define a new point representation for < x, y, z, curvature >
 * 从父类pcl::PointRepresentation中继承
 */
class MyPointRepresentation : public pcl::PointRepresentation<PointNormalT>
{
using pcl::PointRepresentation<PointNormalT>::nr_dimensions_;
public:
    MyPointRepresentation()
    {
        // Define the number of dimensions
        nr_dimensions_ = 4;
    }

    // Override the copyToFloatArray method to define our feature vector
    // 父类中copyToFloatArray是纯虚函数, 需要在子类中被实现
    virtual void copyToFloatArray(const PointNormalT &p, float *out) const
    {
        // < x, y, z, curvature >
        out[0] = p.x;
        out[1] = p.y;
        out[2] = p.z;
        out[3] = p.curvature;
    }
};

/** 
 * @brief Display source and target on the first viewport of the visualizer
 *
 */
void showCloudsLeft(const PointCloud::Ptr cloud_target, const PointCloud::Ptr cloud_source)
{
    p->removePointCloud("vp1_target");
    p->removePointCloud("vp1_source");

    PointCloudColorHandlerCustom<PointT> tgt_h(cloud_target, 0, 255, 0);    // green for target
    PointCloudColorHandlerCustom<PointT> src_h(cloud_source, 255, 0, 0);    // red for source
    p->addPointCloud(cloud_target, tgt_h, "vp1_target", vp_1);
    p->addPointCloud(cloud_source, src_h, "vp1_source", vp_1);

    PCL_INFO("Press q to begin the registration.\n");
    p->spin();
}

/** 
 * @brief Display source and target on the second viewport of the visualizer
 *
 */
void showCloudsRight(const PointCloudWithNormals::Ptr cloud_target, const PointCloudWithNormals::Ptr cloud_source)
{
    p->removePointCloud("source");
    p->removePointCloud("target");

    PointCloudColorHandlerGenericField<PointNormalT> tgt_color_handler(cloud_target, "curvature");
    if (!tgt_color_handler.isCapable())
        PCL_WARN("Cannot create curvature color handler!");

    PointCloudColorHandlerGenericField<PointNormalT> src_color_handler(cloud_source, "curvature");
    if (!src_color_handler.isCapable())
        PCL_WARN("Cannot create curvature color handler!");

    p->addPointCloud(cloud_target, tgt_color_handler, "target", vp_2);
    p->addPointCloud(cloud_source, src_color_handler, "source", vp_2);

    p->spinOnce();
}

/**
 * @brief Load a set of PCD files that we want to register together
 * 
 * @param argc the number of arguments (pass from main ())
 * @param argv the actual command line arguments (pass from main ())
 * @param models the resultant vector of point cloud datasets
 */
void loadData(int argc, char **argv, std::vector<PCD, Eigen::aligned_allocator<PCD>> &models)
{
    std::string extension(".pcd");

    // check pcd file name
    for (int i = 1; i < argc; i++)
    {
        std::string fname = std::string(argv[i]);

        // check filename size
        if (fname.size() <= extension.size())
        {
            continue;
        }

        // convert filename to lowercase
        std::transform(fname.begin(), fname.end(), fname.begin(), (int (*)(int))tolower);

        // check that the argument is a pcd file
        if (fname.compare(fname.size() - extension.size(), extension.size(), extension) == 0)
        {
            // Load the cloud and saves it into the global list of models
            PCD m;
            m.f_name = argv[i];
            pcl::io::loadPCDFile(argv[i], *m.cloud);

            // remove NAN points from the cloud
            std::vector<int> indices;
            pcl::removeNaNFromPointCloud(*m.cloud, *m.cloud, indices);

            models.push_back(m);
        }
    }
}

/**
 * @brief Align a pair of PointCloud datasets and return the result
 * 
 * @param cloud_src the source PointCloud
 * @param cloud_tgt the target PointCloud
 * @param output the resultant aligned source PointCloud
 * @param final_transform the resultant transform between source and target
 * @param downsample 
 */
void pairAlign(
    const PointCloud::Ptr cloud_src, 
    const PointCloud::Ptr cloud_tgt, 
    PointCloud::Ptr output, 
    Eigen::Matrix4f &final_transform, 
    bool downsample = false)
{
    /**
     * @brief Downsample for consistency and speed
     * @note enable this for large datasets
     */
    PointCloud::Ptr src(new PointCloud);
    PointCloud::Ptr tgt(new PointCloud);
    pcl::VoxelGrid<PointT> grid;
    if (downsample)
    {
        grid.setLeafSize(VoxelGridLeafSize, VoxelGridLeafSize, VoxelGridLeafSize);
        grid.setInputCloud(cloud_src);
        grid.filter(*src);

        grid.setInputCloud(cloud_tgt);
        grid.filter(*tgt);
    }
    else
    {
        src = cloud_src;
        tgt = cloud_tgt;
    }

    // Compute surface normals and curvature
    PointCloudWithNormals::Ptr points_with_normals_src(new PointCloudWithNormals);
    PointCloudWithNormals::Ptr points_with_normals_tgt(new PointCloudWithNormals);
    pcl::NormalEstimation<PointT, PointNormalT> norm_est;
    pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>());
    norm_est.setSearchMethod(tree);
    norm_est.setKSearch(NormEstKSearch);

    norm_est.setInputCloud(src);
    norm_est.compute(*points_with_normals_src);
    pcl::copyPointCloud(*src, *points_with_normals_src);

    norm_est.setInputCloud(tgt);
    norm_est.compute(*points_with_normals_tgt);
    pcl::copyPointCloud(*tgt, *points_with_normals_tgt);

    // init custom point representation
    MyPointRepresentation point_representation;
    // weight the 'curvature' dimension so that it is balanced against x, y, and z
    float alpha[4] = {1.0, 1.0, 1.0, 1.0};
    point_representation.setRescaleValues(alpha);

    // ICPNonLinear for registration
    pcl::IterativeClosestPointNonLinear<PointNormalT, PointNormalT> reg;
    reg.setTransformationEpsilon(ICPEPS);
    // Set the maximum distance between two correspondences (src<->tgt) to 10cm
    // Note: adjust this based on the size of your datasets
    reg.setMaxCorrespondenceDistance(ICPMaxDist);
    // Set the point representation
    reg.setPointRepresentation(boost::make_shared<const MyPointRepresentation>(point_representation));
    // set source and target clouds
    reg.setInputSource(points_with_normals_src);
    reg.setInputTarget(points_with_normals_tgt);

    // Run the same optimization in a loop and visualize the results
    Eigen::Matrix4f Ti = Eigen::Matrix4f::Identity(), prev, targetToSource;
    PointCloudWithNormals::Ptr reg_result = points_with_normals_src;
    reg.setMaximumIterations(ICPMaxInnerIteration);
    for (int i = 0; i < ICPMaxExternIteration; ++i)
    {
        PCL_INFO("Iteration Nr. %d.\n", i);

        // save cloud for visualization purpose
        points_with_normals_src = reg_result;

        // ICP Estimate
        reg.setInputSource(points_with_normals_src);
        reg.align(*reg_result);

        // accumulate transformation between each Iteration
        Ti = reg.getFinalTransformation() * Ti;

        // if the difference between this transformation and the previous one
        // is smaller than the threshold, refine the process by reducing
        // the maximal correspondence distance
        if (fabs((reg.getLastIncrementalTransformation() - prev).sum()) < reg.getTransformationEpsilon())
        {
            reg.setMaxCorrespondenceDistance(reg.getMaxCorrespondenceDistance() - 0.001);
        }

        // save current results as prev
        prev = reg.getLastIncrementalTransformation();

        // visualize current state
        showCloudsRight(points_with_normals_tgt, points_with_normals_src);
    }

    // Get the transformation from target to source
    targetToSource = Ti.inverse();

    // Transform target back in source frame 
    pcl::transformPointCloud(*cloud_tgt, *output, targetToSource);

    p->removePointCloud("source");
    p->removePointCloud("target");

    PointCloudColorHandlerCustom<PointT> cloud_tgt_h(output, 0, 255, 0);
    PointCloudColorHandlerCustom<PointT> cloud_src_h(cloud_src, 255, 0, 0);
    p->addPointCloud(output, cloud_tgt_h, "target", vp_2);
    p->addPointCloud(cloud_src, cloud_src_h, "source", vp_2);

    PCL_INFO("Press q to continue the registration.\n");
    p->spin();

    p->removePointCloud("source");
    p->removePointCloud("target");

    // add the source to the transformed target
    *output += *cloud_src;

    final_transform = targetToSource;
}

/**
 * @brief main
 * 
 * @param argc 
 * @param argv 
 * @return int 
 */
int main(int argc, char **argv)
{
    // Load data
    std::vector<PCD, Eigen::aligned_allocator<PCD>> data;
    loadData(argc, argv, data);

    // Check user input
    if (data.empty())
    {
        PCL_ERROR("Syntax is: %s <source.pcd> <target.pcd> [*]", argv[0]);
        PCL_ERROR("[*] - multiple files can be added. The registration results of (i, i+1) will be registered against (i+2), etc");
        return (-1);
    }
    PCL_INFO("Loaded %d datasets.", (int)data.size());

    // Create a PCLVisualizer object
    p = new pcl::visualization::PCLVisualizer(argc, argv, "Pairwise Incremental Registration example");
    p->createViewPort(0.0, 0, 0.5, 1.0, vp_1);
    p->createViewPort(0.5, 0, 1.0, 1.0, vp_2);

    // loop clouds to get transforms
    PointCloud::Ptr result(new PointCloud), source, target;
    Eigen::Matrix4f GlobalTransform = Eigen::Matrix4f::Identity(), pairTransform;
    for (size_t i = 1; i < data.size(); ++i)
    {
        source = data[i - 1].cloud;
        target = data[i].cloud;

        // Add visualization data
        showCloudsLeft(source, target);

        // print source and target info
        PCL_INFO("Aligning %s (%d) with %s (%d).\n", 
            data[i - 1].f_name.c_str(), source->points.size(), 
            data[i].f_name.c_str(), target->points.size());

        // pair align
        PointCloud::Ptr temp(new PointCloud);
        pairAlign(source, target, temp, pairTransform, true);

        // transform current pair into the global transform
        pcl::transformPointCloud(*temp, *result, GlobalTransform);

        // update the global transform
        GlobalTransform = GlobalTransform * pairTransform;

        // save aligned pair, transformed into the first cloud's frame
        std::stringstream ss;
        ss << i << ".pcd";
        pcl::io::savePCDFile(ss.str(), *result, true);
    }
}
